from flask import Flask, jsonify, request, send_from_directory
from flask_cors import CORS
import os
from werkzeug.utils import secure_filename
from datetime import datetime
from image_processing import preprocess_image
from model_inference import run_vehicle_detection, run_plate_detection, run_vehicle_plate_detection, run_multi_vehicle_detection
from ocr_service import recognize_license_plate
import numpy as np
from config import (
    UPLOAD_BASE,
    RAW_IMAGES_DIR,
    PREPROCESSED_IMAGES_DIR,
    INFERENCE_RESULTS_DIR,
    MARKED_VEHICLE_DIR,
    MARKED_PLATE_DIR,
    CROPPED_PLATE_DIR
)
from logger_config import setup_logger

# 创建日志记录器
logger = setup_logger('app')

app = Flask(__name__, 
    static_folder='uploads',  # 将静态文件目录设置为uploads
    static_url_path='/static'  # 设置静态文件URL前缀为/static
)
CORS(app)

@app.route('/')
def index():
    logger.info("访问根路由")
    return jsonify({"message": "车辆识别与车牌检测系统后端服务已启动！"})

@app.route('/api/status')
def status():
    logger.info("检查系统状态")
    return jsonify({
        "status": "运行中",
        "version": "1.0.0",
        "models": ["车辆识别", "车牌检测", "多车检测"],
        "server_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    })

@app.route('/api/upload', methods=['POST'])
def upload_image():
    logger.info("收到图片上传请求")
    try:
        if 'image' not in request.files:
            error_msg = "没有上传文件"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        file = request.files['image']
        if file.filename == '':
            error_msg = "没有选择文件"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        # 生成安全的文件名（使用时间戳）
        timestamp = datetime.now().strftime('%m%d%H%M%S')
        filename = f"{timestamp}_0.jpg"  # 0 表示原始图片
        logger.info(f"生成文件名: {filename}")
        
        # 保存原始文件
        raw_file_path = os.path.join(RAW_IMAGES_DIR, filename)
        file.save(raw_file_path)
        logger.info(f"原始图片已保存至: {raw_file_path}")
        
        # 获取原始文件大小
        file_size = os.path.getsize(raw_file_path)
        logger.debug(f"文件大小: {file_size} 字节")
        
        # 预处理图像
        processed_path = preprocess_image(raw_file_path, PREPROCESSED_IMAGES_DIR)
        logger.info(f"图片预处理完成: {processed_path}")
        
        # 获取预处理后图片的相对路径
        processed_filename = os.path.basename(processed_path)
        preprocessed_image_url = f"/static/{os.path.relpath(processed_path, UPLOAD_BASE)}"
        
        response = {
            "message": "图片上传和预处理成功",
            "original_filename": filename,
            "file_size": file_size,
            "preprocessed_image_url": preprocessed_image_url,
            "preprocessed_image_path": processed_path
        }
        logger.info(f"返回预处理结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"处理失败: {str(e)}"
        logger.error(error_msg)
        # 如果处理过程中出现错误，尝试清理已上传的文件
        if 'raw_file_path' in locals() and os.path.exists(raw_file_path):
            os.remove(raw_file_path)
            logger.info(f"已清理临时文件: {raw_file_path}")
        return jsonify({"error": error_msg}), 500

@app.route('/api/vehicle_detection', methods=['POST'])
def vehicle_detection():
    logger.info("收到车辆检测请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"车辆检测请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行车辆检测
        marked_image_path, detections = run_vehicle_detection(
            preprocessed_image_path,
            "vehicle_detection",
            MARKED_VEHICLE_DIR
        )
        
        # 计算检测结果的平均置信度
        if detections:
            overall_accuracy = np.mean([det['confidence'] for det in detections])
        else:
            overall_accuracy = 0.0
        
        # 生成标注后图片的URL
        marked_image_url = f"/static/{os.path.relpath(marked_image_path, UPLOAD_BASE)}"
        
        response = {
            "message": "车辆检测成功",
            "marked_image_url": marked_image_url,
            "detection_results": detections,
            "overall_accuracy": round(float(overall_accuracy), 2),
            "detection_count": len(detections)
        }
        logger.info(f"返回车辆检测结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"车辆检测失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/license_plate_recognition', methods=['POST'])
def license_plate_recognition():
    logger.info("收到车牌检测请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"车牌检测请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行车牌检测
        marked_image_path, cropped_plate_path, ocr_res_img_path, plate_text, detections = run_plate_detection(
            preprocessed_image_path,
            "plate_detection",
            MARKED_PLATE_DIR
        )
        
        # 计算检测结果的平均置信度
        if detections:
            overall_accuracy = np.mean([det['confidence'] for det in detections])
        else:
            overall_accuracy = 0.0
        
        # 生成标注后图片的URL
        marked_image_url = f"/static/{os.path.relpath(marked_image_path, UPLOAD_BASE)}"
        
        # 生成切割后车牌图片的URL（如果有）
        cropped_plate_url = None
        if cropped_plate_path:
            cropped_plate_url = f"/static/{os.path.relpath(cropped_plate_path, UPLOAD_BASE)}"
        
        # 生成OCR结果图片的URL（如果有）
        ocr_res_img_url = None
        if ocr_res_img_path:
            ocr_res_img_url = f"/static/{os.path.relpath(ocr_res_img_path, UPLOAD_BASE)}"
        
        response = {
            "message": "车牌检测成功",
            "marked_image_url": marked_image_url,
            "detection_results": detections,
            "overall_accuracy": round(float(overall_accuracy), 2),
            "detection_count": len(detections)
        }
        
        # 如果有切割后的车牌图片，添加到响应中
        if cropped_plate_url:
            response["cropped_plate_url"] = cropped_plate_url
        
        # 如果有OCR识别结果，添加到响应中
        if plate_text:
            response["plate_text"] = plate_text
        
        # 如果有OCR结果图片，添加到响应中
        if ocr_res_img_url:
            response["ocr_res_img_url"] = ocr_res_img_url
        
        logger.info(f"返回车牌检测结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"车牌检测失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/full_process', methods=['POST'])
def full_process():
    logger.info("收到完整处理请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"完整处理请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行车辆检测
        vehicle_marked_image_path, vehicle_detections = run_vehicle_detection(
            preprocessed_image_path,
            "vehicle_detection",
            MARKED_VEHICLE_DIR
        )
        
        # 执行车牌检测
        plate_marked_image_path, cropped_plate_path, ocr_res_img_path, plate_text, plate_detections = run_plate_detection(
            preprocessed_image_path,
            "plate_detection",
            MARKED_PLATE_DIR
        )
        
        # 计算车辆检测的平均置信度
        if vehicle_detections:
            vehicle_accuracy = np.mean([det['confidence'] for det in vehicle_detections])
        else:
            vehicle_accuracy = 0.0
        
        # 计算车牌检测的平均置信度
        if plate_detections:
            plate_accuracy = np.mean([det['confidence'] for det in plate_detections])
        else:
            plate_accuracy = 0.0
        
        # 计算整体的平均置信度
        overall_accuracy = (vehicle_accuracy + plate_accuracy) / 2 if vehicle_detections or plate_detections else 0.0
        
        # 生成标注后图片的URL
        vehicle_image_url = f"/static/{os.path.relpath(vehicle_marked_image_path, UPLOAD_BASE)}"
        plate_image_url = f"/static/{os.path.relpath(plate_marked_image_path, UPLOAD_BASE)}"
        
        # 生成切割后车牌图片的URL（如果有）
        cropped_plate_url = None
        if cropped_plate_path:
            cropped_plate_url = f"/static/{os.path.relpath(cropped_plate_path, UPLOAD_BASE)}"
        
        # 生成OCR结果图片的URL（如果有）
        ocr_res_img_url = None
        if ocr_res_img_path:
            ocr_res_img_url = f"/static/{os.path.relpath(ocr_res_img_path, UPLOAD_BASE)}"
        
        # 构建响应
        response = {
            "message": "完整处理成功",
            "vehicle_detection": {
                "marked_image_url": vehicle_image_url,
                "detection_results": vehicle_detections,
                "accuracy": round(float(vehicle_accuracy), 2),
                "detection_count": len(vehicle_detections)
            },
            "license_plate_recognition": {
                "marked_image_url": plate_image_url,
                "detection_results": plate_detections,
                "accuracy": round(float(plate_accuracy), 2),
                "detection_count": len(plate_detections)
            },
            "overall_accuracy": round(float(overall_accuracy), 2)
        }
        
        # 如果有切割后的车牌图片，添加到响应中
        if cropped_plate_url:
            response["license_plate_recognition"]["cropped_plate_url"] = cropped_plate_url
        
        # 如果有OCR识别结果，添加到响应中
        if plate_text:
            response["license_plate_recognition"]["plate_text"] = plate_text
        
        # 如果有OCR结果图片，添加到响应中
        if ocr_res_img_url:
            response["license_plate_recognition"]["ocr_res_img_url"] = ocr_res_img_url
        
        logger.info(f"返回完整处理结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"完整处理失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/ocr_recognition', methods=['POST'])
def ocr_recognition():
    logger.info("收到OCR识别请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        cropped_plate_path = data.get('cropped_plate_path')
        if not cropped_plate_path:
            error_msg = "缺少车牌图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"OCR识别请求参数: 图片路径={cropped_plate_path}")
        
        # 检查文件是否存在
        if not os.path.exists(cropped_plate_path):
            error_msg = "车牌图片不存在"
            logger.error(f"{error_msg}: {cropped_plate_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行OCR识别
        ocr_res_img_path, plate_text, ocr_confidence = recognize_license_plate(
            cropped_plate_path,
            os.path.dirname(cropped_plate_path)
        )
        
        # 生成OCR结果图片的URL
        ocr_res_img_url = f"/static/{os.path.relpath(ocr_res_img_path, UPLOAD_BASE)}"
        
        response = {
            "message": "OCR识别成功",
            "plate_text": plate_text,
            "ocr_confidence": round(float(ocr_confidence), 2),
            "ocr_res_img_url": ocr_res_img_url
        }
        
        logger.info(f"返回OCR识别结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"OCR识别失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/vehicle_plate_detection', methods=['POST'])
def vehicle_plate_detection():
    logger.info("收到车辆-车牌联合检测请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"车辆-车牌联合检测请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行车辆-车牌联合检测
        marked_image_path, detections, plate_text = run_vehicle_plate_detection(
            preprocessed_image_path,
            INFERENCE_RESULTS_DIR
        )
        
        # 分离车辆和车牌检测结果
        vehicle_detections = [det for det in detections if det['type'] == '车辆']
        plate_detections = [det for det in detections if det['type'] == '车牌']
        
        # 计算车辆检测的平均置信度
        if vehicle_detections:
            vehicle_accuracy = np.mean([det['confidence'] for det in vehicle_detections])
        else:
            vehicle_accuracy = 0.0
        
        # 计算车牌检测的平均置信度
        if plate_detections:
            plate_accuracy = np.mean([det['confidence'] for det in plate_detections])
        else:
            plate_accuracy = 0.0
        
        # 计算整体的平均置信度
        overall_accuracy = (vehicle_accuracy + plate_accuracy) / 2 if vehicle_detections or plate_detections else 0.0
        
        # 生成标注后图片的URL
        marked_image_url = f"/static/{os.path.relpath(marked_image_path, UPLOAD_BASE)}"
        
        # 构建响应
        response = {
            "message": "车辆-车牌联合检测成功",
            "marked_image_url": marked_image_url,
            "vehicle_detections": vehicle_detections,
            "plate_detections": plate_detections,
            "vehicle_accuracy": round(float(vehicle_accuracy), 2),
            "plate_accuracy": round(float(plate_accuracy), 2),
            "overall_accuracy": round(float(overall_accuracy), 2),
            "vehicle_count": len(vehicle_detections),
            "plate_count": len(plate_detections)
        }
        
        # 如果有OCR识别结果，添加到响应中
        if plate_text:
            response["plate_text"] = plate_text
        
        logger.info(f"返回车辆-车牌联合检测结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"车辆-车牌联合检测失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/multi_vehicle_detection', methods=['POST'])
def multi_vehicle_detection():
    logger.info("收到多车检测请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"多车检测请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行多车检测
        marked_image_path, detections = run_multi_vehicle_detection(
            preprocessed_image_path,
            "multi_vehicle_detection",
            MARKED_VEHICLE_DIR
        )
        
        # 计算检测结果的平均置信度
        if detections:
            overall_accuracy = np.mean([det['confidence'] for det in detections])
        else:
            overall_accuracy = 0.0
        
        # 生成标注后图片的URL
        marked_image_url = f"/static/{os.path.relpath(marked_image_path, UPLOAD_BASE)}"
        
        response = {
            "message": "多车检测成功",
            "marked_image_url": marked_image_url,
            "detection_results": detections,
            "overall_accuracy": round(float(overall_accuracy), 2),
            "detection_count": len(detections)
        }
        logger.info(f"返回多车检测结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"多车检测失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

@app.route('/api/multi_vehicle_full_process', methods=['POST'])
def multi_vehicle_full_process():
    logger.info("收到多车完整处理请求")
    try:
        # 获取请求数据
        data = request.get_json()
        if not data:
            error_msg = "无效的请求数据"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        preprocessed_image_path = data.get('preprocessed_image_path')
        if not preprocessed_image_path:
            error_msg = "缺少预处理图像路径"
            logger.error(error_msg)
            return jsonify({"error": error_msg}), 400
        
        logger.info(f"多车完整处理请求参数: 图片路径={preprocessed_image_path}")
        
        # 检查文件是否存在
        if not os.path.exists(preprocessed_image_path):
            error_msg = "预处理图片不存在"
            logger.error(f"{error_msg}: {preprocessed_image_path}")
            return jsonify({"error": error_msg}), 404
        
        # 执行多车检测
        vehicle_marked_image_path, vehicle_detections = run_multi_vehicle_detection(
            preprocessed_image_path,
            "multi_vehicle_detection",
            MARKED_VEHICLE_DIR
        )
        
        # 执行车牌检测
        plate_marked_image_path, cropped_plate_path, ocr_res_img_path, plate_text, plate_detections = run_plate_detection(
            preprocessed_image_path,
            "plate_detection",
            MARKED_PLATE_DIR
        )
        
        # 计算多车检测的平均置信度
        if vehicle_detections:
            vehicle_accuracy = np.mean([det['confidence'] for det in vehicle_detections])
        else:
            vehicle_accuracy = 0.0
        
        # 计算车牌检测的平均置信度
        if plate_detections:
            plate_accuracy = np.mean([det['confidence'] for det in plate_detections])
        else:
            plate_accuracy = 0.0
        
        # 计算整体的平均置信度
        overall_accuracy = (vehicle_accuracy + plate_accuracy) / 2 if vehicle_detections or plate_detections else 0.0
        
        # 生成标注后图片的URL
        vehicle_image_url = f"/static/{os.path.relpath(vehicle_marked_image_path, UPLOAD_BASE)}"
        plate_image_url = f"/static/{os.path.relpath(plate_marked_image_path, UPLOAD_BASE)}"
        
        # 生成切割后车牌图片的URL（如果有）
        cropped_plate_url = None
        if cropped_plate_path:
            cropped_plate_url = f"/static/{os.path.relpath(cropped_plate_path, UPLOAD_BASE)}"
        
        # 生成OCR结果图片的URL（如果有）
        ocr_res_img_url = None
        if ocr_res_img_path:
            ocr_res_img_url = f"/static/{os.path.relpath(ocr_res_img_path, UPLOAD_BASE)}"
        
        # 构建响应
        response = {
            "message": "多车完整处理成功",
            "vehicle_detection": {
                "marked_image_url": vehicle_image_url,
                "detection_results": vehicle_detections,
                "accuracy": round(float(vehicle_accuracy), 2),
                "detection_count": len(vehicle_detections)
            },
            "license_plate_recognition": {
                "marked_image_url": plate_image_url,
                "detection_results": plate_detections,
                "accuracy": round(float(plate_accuracy), 2),
                "detection_count": len(plate_detections)
            },
            "overall_accuracy": round(float(overall_accuracy), 2)
        }
        
        # 如果有切割后的车牌图片，添加到响应中
        if cropped_plate_url:
            response["license_plate_recognition"]["cropped_plate_url"] = cropped_plate_url
        
        # 如果有OCR识别结果，添加到响应中
        if plate_text:
            response["license_plate_recognition"]["plate_text"] = plate_text
        
        # 如果有OCR结果图片，添加到响应中
        if ocr_res_img_url:
            response["license_plate_recognition"]["ocr_res_img_url"] = ocr_res_img_url
        
        logger.info(f"返回多车完整处理结果: {response}")
        return jsonify(response)
        
    except Exception as e:
        error_msg = f"多车完整处理失败: {str(e)}"
        logger.error(error_msg)
        return jsonify({"error": error_msg}), 500

if __name__ == '__main__':
    logger.info("启动 Flask 应用")
    # 确保上传目录存在
    for dir_path in [RAW_IMAGES_DIR, PREPROCESSED_IMAGES_DIR, MARKED_VEHICLE_DIR, MARKED_PLATE_DIR, CROPPED_PLATE_DIR]:
        os.makedirs(dir_path, exist_ok=True)
        logger.info(f"确保目录存在: {dir_path}")
    
    app.run(host='0.0.0.0', port=5000, debug=True)